import asyncio
from datetime import datetime
from typing import AsyncGenerator, Dict
import json
import os
from ...strategy_agent.client import StrategyClient
from ...data_agent.client import DataClient

class BacktestRunner:
    def __init__(self):
        self.percent_complete = 0
        self.cache_interval = 1000  # 每处理1000条缓存一次
        self.cache_counter = 0
        self.cache_path = f"/tmp/backtest_cache_{datetime.now().timestamp()}.json"
        self.last_flush = datetime.now()

    async def initialize(self, strategy_id: str, start_time: datetime, end_time: datetime):
        """初始化回测参数并加载缓存"""
        self.strategy_id = strategy_id
        self.start_time = start_time
        self.end_time = end_time
        self.cache_counter = 0
        
        # 加载已有缓存
        if os.path.exists(self.cache_path):
            with open(self.cache_path, 'r') as f:
                cache = json.load(f)
                self.strategy = StrategyClient.from_cache(cache['strategy'])
                self.percent_complete = cache['progress']
            return

        # 初始化新回测
        self.strategy = await StrategyClient.load_strategy(strategy_id)
        self.data_stream = DataClient.stream_ticks(
            symbol=self.strategy.symbol,
            start=start_time,
            end=end_time
        )

    async def _save_cache(self):
        """保存中间结果到缓存文件"""
        cache_data = {
            'strategy': self.strategy.to_dict(),
            'progress': self.percent_complete,
            'timestamp': datetime.now().isoformat()
        }
        with open(self.cache_path, 'w') as f:
            json.dump(cache_data, f)
        self.last_flush = datetime.now()

    async def execute(self) -> AsyncGenerator[dict, None]:
        """执行回测并流式返回进度（带缓存机制）"""
        total_seconds = (self.end_time - self.start_time).total_seconds()
        start_timestamp = self.start_time.timestamp()
        
        async for tick in self.data_stream:
            self.cache_counter += 1
            processed = await self.strategy.process(tick)
            
            # 缓存中间结果
            if self.cache_counter % self.cache_interval == 0:
                await self._save_cache()
            
            # 计算进度
            elapsed = tick.timestamp.timestamp() - start_timestamp
            self.percent_complete = min(99, int(elapsed / total_seconds * 100))
            
            yield {
                "percent": self.percent_complete,
                "metrics": self._calculate_metrics(),
                "equity_curve": self.strategy.portfolio.equity_history[-100:]
            }
            
        # 最终完成并清理缓存
        await self._save_cache()
        os.remove(self.cache_path)
        yield {
            "percent": 100,
            "metrics": self._calculate_metrics(final=True),
            "equity_curve": self.strategy.portfolio.equity_history
        }

    # 保持其他方法不变...